function [RegionSet, SampleSet, RegionSampleId] = PRS_MMO( Problem, AlgorithmM )
% PRS_MMO: Partition-Based Random Search for Multimodal Optimization for minimization problems
figureNum=1;
%INPUT
%   Problem.
%       Dimension: scalar, the dimension of the decision variables
%       Domain: 1 -by- ~ vector, features of the domain
%       Partition: function handle of partitioning strategy
%           [RegionNum, Region, RegionSampleId] = Partition(Region, SampleSet)
%               RegionNum: scalar, number of new partitioned regions
%               Region = [partition depth, partitionable, features ...]
%                   partition depth = 0 for the whole domain
%                   partitionable: scalar, 1: partitionable, 0: non-partitionable
%               RegionSampleId: ~-by-1 cell, the sample id that belongs to each region
%               SampleSet: [SampleID, objective value, x1, x2, ...]
%       Sampling: function handle of sampling in a particular region
%           [Sample] = Sampling(n, region)
%               n: number of samples
%               region: 1 -by- ~ vector, [features...]
%               Sample: n -by- (1+d) matrix, [objective value, x1, x2, ...]
%   AlgorithmM.
%       n0: scalar, first stage sampling size
%       QuantileLevel: scalar, <0.5, the quantile level of promisinhg region definition
%       NewBudget: scalar, added budget size at each stage
%       MaxSampleSize: scalar, maximal sample size in one region
%       StopCriteria: [X,Y]
%           [1,Y]:reach maximum budget Y
%           [2,Y]:reach acceptable optimum Y

%OUTPUT
%   RegionSet: ~ -by- (2+~) matrix, the whole subregion set
%         [partition depth, partitionable, features ...]
%   SampleSet: ~ -by- (1+d) matrix, the original sample set
%         [objective value, x1, x2, ...]
%   RegionSampleId: ~ -by- ~ cell, the id of the samples in each region

%% Initialization
% PROBLEM PARAMETERS
d = Problem.Dimension;
domain = [0, 1, Problem.Domain]; % [partition depth, partitionable, features ...]
Partition = Problem.Partition;
Sampling = Problem.Sampling;
% ALGORITHM PARAMETERS
n0 = AlgorithmM.n0;
alpha = AlgorithmM.QuantileLevel;
Delta = AlgorithmM.NewBudget;
MaxSampleSize = AlgorithmM.MaxSampleSize;
StopCriteria = AlgorithmM.StopCriteria;
N = (StopCriteria(1)==1)*StopCriteria(2) + (StopCriteria(1)~=1)*1000000; % the maximal budget size

%% Optimization
% PRE-PROCESSING
RegionSet = zeros(round(N/n0), length(domain));
RegionSet(1,:) = domain;
RegionsNum = 1; % the number of exist regions
CurrentMaxPartition = 0;  % the maximum depth of partitioning at current iteration
GroupMean = zeros(round(N/n0), 1); % group sample mean
GroupStd = zeros(round(N/n0), 1); % group sample std
GroupN1 = zeros(round(N/n0), 1); % real group sample size
GroupN1_mod = zeros(round(N/n0), 1); % modified group sample size
RegionSampleId = cell(round(N/n0), 1); % id of group samples
SampleSet = zeros(N, 2+d);  % [SampleID,objective value,x1,x2,...]
SampleSet(:,1) = 1:size(SampleSet,1);
N0 = 0; % budget allocated

% START OPTIMIZATION
CurrentOpt = inf;  % the current optimal value
PartitionList = zeros(Delta,1);
PartitionList(1) = 1;
PartitionListNum = 1;
% subsequent iterations
while N0 < N...
    && (StopCriteria(1)~=2 || CurrentOpt>StopCriteria(2))...
%% partition partitionable regions
    i = 1;
    while i <= PartitionListNum
        iPartitionRegion = PartitionList(i);
        [NewRegionNum, NewRegion, NewRegionSampleId]...
            = Partition(RegionSet(iPartitionRegion,:), SampleSet(RegionSampleId{iPartitionRegion},:));
        NewRegionId = [iPartitionRegion,(RegionsNum+1):(RegionsNum+NewRegionNum-1)]; % id of new-partitioned region
        RegionSet(NewRegionId,:) = NewRegion;
        RegionSampleId(NewRegionId) = NewRegionSampleId;
        RegionsNum = RegionsNum+NewRegionNum-1; % update total region number
        % update partition state
        PartitionDepth_temp = max(NewRegion(:,1));
        if PartitionDepth_temp>CurrentMaxPartition
            CurrentMaxPartition = PartitionDepth_temp;
        end
        % first stage sampling and group information updating
        for j = NewRegionId
            GroupN1(j) = length(RegionSampleId{j});
            if GroupN1(j)>=MaxSampleSize && RegionSet(j,2)==1
                PartitionListNum = PartitionListNum+1;
                PartitionList(PartitionListNum) = j;
            else
                % first stage sampling
                if RegionSet(j,2)==0 % non-partitionable
                    Add_n = min(1-GroupN1(j),N-N0);
                else
                    Add_n = min(n0-GroupN1(j),N-N0);
                end
                if Add_n>0
                    NewSampleId = ((N0+1):(N0+Add_n))';
                    SampleSet(NewSampleId,2:end) = Sampling(Add_n,RegionSet(j,3:end));
                    CurrentOpt = min([CurrentOpt;SampleSet(NewSampleId,2)]);
                    N0 = N0+Add_n;
                    RegionSampleId{j} = [RegionSampleId{j};NewSampleId];
                    GroupN1(j) = GroupN1(j)+Add_n;
                end
                % update group information
                Samplevalue_temp = SampleSet(RegionSampleId{j},2);
                GroupMean(j) = mean(Samplevalue_temp);
                GroupStd(j) = sqrt(var(Samplevalue_temp));
            end
        end
        i = i+1;
    end
    PartitionListNum = 0;
    
%% iteration figures
if (N0/500)>=figureNum
    figure()
    hold on
    rectangle('Position',[Problem.Domain([1,3]),Problem.Domain([2,4])-Problem.Domain([1,3])])
    scatter(SampleSet(:,3),SampleSet(:,4),2,SampleSet(:,2),'filled');
    for i=1:length(RegionSet)
        rectangle('Position',[RegionSet(i,[3,5]),RegionSet(i,[4,6])-RegionSet(i,[3,5])])
    end
    title('','Interpreter','latex','String',['$N=',num2str(N0),'$'])
    xlabel('','Interpreter','latex','String','$x_1$')
    ylabel('','Interpreter','latex','String','$x_2$')
    set(gca,'Fontname','Times New Roman','FontSize',16);
    hold off
    figureNum=figureNum+1;
end

%% modify the BAQM equations
    ModConst = RegionSet(1:RegionsNum,1)./CurrentMaxPartition;
    GroupN1_mod(1:RegionsNum)=max(2,round(GroupN1(1:RegionsNum).*ModConst));
%% additional budget
    Add_n = min(N-N0,Delta);
    n2 = BAQM(Add_n,alpha,GroupMean(1:RegionsNum)',GroupStd(1:RegionsNum)',GroupN1_mod(1:RegionsNum)', RegionSet(1:RegionsNum,2)');
    RegionId_temp = find(n2>0);
    for i = RegionId_temp
        NewSampleId = ((N0+1):(N0+n2(i)))';
        SampleSet(NewSampleId,2:end) = Sampling(n2(i),RegionSet(i,3:end));
        CurrentOpt = min([CurrentOpt;SampleSet(NewSampleId,2)]);
        N0 = N0+n2(i);
        RegionSampleId{i} = [RegionSampleId{i};NewSampleId];
        mean_temp = GroupMean(i);
        GroupMean(i) = (GroupMean(i)*GroupN1(i)+sum(SampleSet(NewSampleId,2)))/(GroupN1(i)+n2(i));
        GroupStd(i) = sqrt((GroupStd(i)^2*(GroupN1(i)-1)...
            +(GroupMean(i)-mean_temp)^2*GroupN1(i)...
            +sum((SampleSet(NewSampleId,2)-GroupMean(i)).^2))/(GroupN1(i)+n2(i)-1));
        GroupN1(i) = GroupN1(i)+n2(i);
        if GroupN1(i)>=MaxSampleSize && RegionSet(i,2)==1
            PartitionListNum = PartitionListNum+1;
            PartitionList(PartitionListNum) = i;
        end
    end
end
RegionSet((RegionsNum+1):end, :) = [];
SampleSet((N0+1):end, :) = [];
SampleSet(:,1)=[];
RegionSampleId((RegionsNum+1):end, :) = [];

%% Plots for 2 dimension box constraint problems
% if d==2
%     figure()
%     hold on
%     for i=1:RegionsNum
%         rectangle('Position',[RegionSet(i,[3,5]),RegionSet(i,[4,6])-RegionSet(i,[3,5])])
%     end
%     scatter(SampleSet(:,2),SampleSet(:,3),2,SampleSet(:,1),'filled');
%     hold off
% end

%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% FUNCTIONS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function [ n2 ] = BAQM( N_new, Percentage, mu_hat, sigma_hat, n1, partitonable )
    %BAQM Budget Allocation for Quantile Minimization
    %INPUT
    %   N_new: the new simulation budget
    %   Percentage: the percentage of good points of interest
    %   mu_hat: 1-by-K vector, sample mean of each group
    %   sigma_hat: 1-by-K vector, sample standard deviation of each group
    %   n1: 1-by-K vector, number of observations in each group at stage one
    %   partitonable: 1-by-K vector, 1: the region is partitionable
    %OUTPUT
    %   n2: 1-by-K vector, additional budgets need to be allocated to each group
        zero_tolerance=1e-16;
        y0=mu_hat+sigma_hat.*norminv(Percentage,0,1);
        [y,b]=min(y0); % threshold and the best group
        %% from here only partitionable regions are considered
        partitonable_temp=partitonable;
        partitonable_temp(b)=0;
        partitionableRegionId=find(partitonable_temp==1);
        partitionableRegionId=[b,partitionableRegionId];
        K=length(partitionableRegionId); % number of groups
        n1_mod=n1(partitionableRegionId);
        cv_hat=sigma_hat(partitionableRegionId)./(mu_hat(partitionableRegionId)-y);
        if sigma_hat(b)<zero_tolerance
            cv_hat(1)=-1/norminv(Percentage,0,1); % if the best group has 0 variance
        end
        % the ration between group k and group b
        Wk=1./(1+1./(cv_hat.^2)-1./n1_mod);
        Ckb=Wk./Wk(1);
        P_kb=fcdf(Ckb,n1_mod-1,ones(1,K).*n1_mod(1)-1); % the probability that cv_k>cv_b
%         if partitonable(b)==0
%             P_kb=fcdf(Ckb,n1_mod-1,ones(1,K).*n1_mod(1)-1); % the probability that cv_k>cv_b
%         else
%             P_kb=chi2cdf(Ckb.*(n1_mod-1),n1_mod-1);
%         end
        P_hat=P_kb./(1-P_kb); % Ni / Nb
        if partitonable(b)==0
            partitionableRegionId=partitionableRegionId(2:end);
            n1_mod=n1_mod(2:end);
            P_hat=P_hat(2:end);
            K=K-1;
            if sum(P_hat)==0
                P_hat = ones(1,K);
            end
        end
        %% new budget allcation
        Ntot=sum(n1_mod)+N_new;
        n=P_hat./sum(P_hat).*Ntot; %total budgets allocation
        % additional budgets allocation
        n_delta=n-n1_mod;
        [~,n_delta_index]=sort(n_delta,'descend'); % ranking the groups
        n2_mod=zeros(1,K);
        N_temp1=N_new;
        for k=1:K
            n2_mod(n_delta_index(k))=max(0,min(N_temp1,round(n_delta(n_delta_index(k)))));
            N_temp1=N_temp1-n2_mod(n_delta_index(k));
            if N_temp1==0
                break
            end
        end
        % approximation caused by rounding to integer
        n_delta=n-n1_mod-n2_mod;
        [~,n_delta_index]=sort(n_delta,'descend'); % ranking the groups
        k=1;
        while sum(n2_mod)<(Ntot-sum(n1_mod))
            n2_mod(n_delta_index(k))=n2_mod(n_delta_index(k))+1;
            k=k+1;
        end
        n2=zeros(1,length(mu_hat));
        n2(partitionableRegionId)=n2_mod;
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
end

